# Author: Cascade
# Date: 2024/06/24
# Description: CosyVoice 2 TTS Service Implementation

import os
import logging
import requests
import json
import time
from typing import List, Dict, Any, Optional, Tuple
from pathlib import Path
import base64
import wave
import io

from pptflow.tts.tts_service import TtsService
from pptflow.config.setting import Setting
from pptflow.utils import mylogger

logger = mylogger.get_logger(__name__)

class CosyVoiceTtsService(TtsService):
    """
    CosyVoice 2 TTS Service Implementation
    Enhanced version with support for CosyVoice 2 API features
    """
    
    def __init__(self):
        # Base configuration
        self.base_url = os.getenv("COSYVOICE_API", "https://api.cosyvoice.ai/v2").rstrip('/')
        self.api_key = os.getenv("COSYVOICE_API_KEY", "")
        
        # Voice settings with enhanced defaults for CosyVoice 2
        self.default_voice = os.getenv("COSYVOICE_DEFAULT_VOICE", "zh-CN-YunxiNeural")
        self.default_speed = float(os.getenv("COSYVOICE_SPEED", "1.0"))
        self.default_pitch = float(os.getenv("COSYVOICE_PITCH", "0.0"))
        self.default_style = os.getenv("COSYVOICE_STYLE", "neutral")
        self.default_emotion = os.getenv("COSYVOICE_EMOTION", "neutral")
        self.sample_rate = int(os.getenv("COSYVOICE_SAMPLE_RATE", "24000"))
        
        # Audio format settings
        self.audio_format = os.getenv("COSYVOICE_AUDIO_FORMAT", "wav")
        self.bit_depth = int(os.getenv("COSYVOICE_BIT_DEPTH", "16"))
        
        # Connection settings
        self.timeout = int(os.getenv("COSYVOICE_TIMEOUT", "30"))
        self.max_retries = int(os.getenv("COSYVOICE_MAX_RETRIES", "3"))
        self.retry_delay = int(os.getenv("COSYVOICE_RETRY_DELAY", "1"))
        
        # Cache for voice list to avoid repeated API calls
        self._voice_list_cache = None
        self._voice_list_last_updated = 0
        self._voice_cache_ttl = 3600  # 1 hour cache
        
        logger.info(f"Initialized CosyVoice 2 TTS Service with {self.base_url}")

    def _get_auth_headers(self) -> Dict[str, str]:
        """Get authentication headers with API key if available"""
        headers = {
            "Content-Type": "application/json",
            "Accept": "application/json",
        }
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        return headers
        
    def _make_request(self, method: str, endpoint: str, **kwargs) -> requests.Response:
        """Make an HTTP request with retry logic and error handling"""
        url = f"{self.base_url}/{endpoint.lstrip('/')}"
        headers = self._get_auth_headers()
        
        # Update headers with any provided in kwargs
        if 'headers' in kwargs:
            headers.update(kwargs['headers'])
            del kwargs['headers']
            
        for attempt in range(self.max_retries + 1):
            try:
                response = requests.request(
                    method=method,
                    url=url,
                    headers=headers,
                    timeout=self.timeout,
                    **kwargs
                )
                
                # If rate limited, wait and retry
                if response.status_code == 429:
                    retry_after = int(response.headers.get('Retry-After', self.retry_delay))
                    logger.warning(f"Rate limited. Waiting {retry_after} seconds before retry...")
                    time.sleep(retry_after)
                    continue
                    
                # For other client errors, don't retry
                if 400 <= response.status_code < 500 and response.status_code != 429:
                    error_msg = f"Client error {response.status_code}"
                    try:
                        error_data = response.json()
                        error_msg = error_data.get('error', {}).get('message', str(error_data))
                    except:
                        error_msg = response.text or error_msg
                    logger.error(f"{error_msg}")
                    response.raise_for_status()
                    
                # For server errors, retry
                if response.status_code >= 500:
                    raise requests.exceptions.HTTPError(f"Server error: {response.status_code}")
                    
                return response
                
            except (requests.exceptions.RequestException, requests.exceptions.HTTPError) as e:
                if attempt == self.max_retries:
                    logger.error(f"Request failed after {self.max_retries} attempts: {str(e)}")
                    raise
                wait_time = self.retry_delay * (2 ** attempt)  # Exponential backoff
                logger.warning(f"Attempt {attempt + 1} failed. Retrying in {wait_time} seconds...")
                time.sleep(wait_time)
    
    async def tts(self, text: str, output_audio_filename: str, setting: Setting):
        """
        Convert text to speech using CosyVoice 2
        
        Args:
            text: The text to convert to speech
            output_audio_filename: Path to save the generated audio
            setting: Application settings
        """
        try:
            if not text.strip():
                logger.warning("Empty text provided for TTS")
                return
            
            # Prepare the request payload for CosyVoice 2 API
            payload = {
                "input": {
                    "text": text
                },
                "voice": {
                    "name": self.default_voice,
                    "style": self.default_style,
                    "emotion": self.default_emotion
                },
                "audio_config": {
                    "audio_encoding": self.audio_format.upper(),
                    "sample_rate_hertz": self.sample_rate,
                    "pitch": self.default_pitch,
                    "speaking_rate": self.default_speed,
                    "bit_depth": self.bit_depth
                }
            }
            
            # Make the API request with retry logic
            response = self._make_request(
                method="POST",
                endpoint="/synthesize",
                json=payload
            )
            
            # Handle the response
            if response.status_code == 200:
                # Handle different response formats
                content_type = response.headers.get('Content-Type', '')
                if 'application/json' in content_type:
                    # If response is JSON with base64 audio
                    result = response.json()
                    if 'audioContent' in result:
                        audio_data = base64.b64decode(result['audioContent'])
                    else:
                        raise ValueError("Invalid response format: missing audioContent")
                else:
                    # If response is raw audio
                    audio_data = response.content
                
                # Save the audio file
                output_path = Path(output_audio_filename)
                output_path.parent.mkdir(parents=True, exist_ok=True)
                
                with open(output_path, 'wb') as f:
                    f.write(audio_data)
                    
                logger.info(f"TTS audio saved to {output_audio_filename}")
                
                # Convert to WAV if needed (for compatibility)
                if output_path.suffix.lower() == '.wav' and not self._is_wav_file(audio_data):
                    self._convert_to_wav(str(output_path), audio_data)
                    
            else:
                error_msg = f"CosyVoice API error: {response.status_code}"
                try:
                    error_data = response.json()
                    error_msg = error_data.get('error', {}).get('message', str(error_data))
                except:
                    error_msg = response.text or error_msg
                logger.error(error_msg)
                raise Exception(f"CosyVoice API error: {error_msg}")
            
        except Exception as e:
            logger.error(f"Error in CosyVoice 2 TTS: {str(e)}")
            raise
    
    def _is_wav_file(self, data: bytes) -> bool:
        """Check if the binary data is a valid WAV file"""
        try:
            with io.BytesIO(data) as f:
                return f.read(4) == b'RIFF'
        except:
            return False
    
    def _convert_to_wav(self, output_path: str, audio_data: bytes):
        """Convert raw audio data to WAV format"""
        try:
            import numpy as np
            import soundfile as sf
            
            # Convert raw bytes to numpy array based on bit depth
            if self.bit_depth == 16:
                audio_array = np.frombuffer(audio_data, dtype=np.int16)
            elif self.bit_depth == 32:
                audio_array = np.frombuffer(audio_data, dtype=np.int32)
            else:
                audio_array = np.frombuffer(audio_data, dtype=np.int16)  # Default to 16-bit
            
            # Save as WAV file
            sf.write(output_path, audio_array, self.sample_rate)
            logger.debug(f"Converted audio to WAV format: {output_path}")
            
        except Exception as e:
            logger.warning(f"Failed to convert audio to WAV: {str(e)}")
            # If conversion fails, just write the original data
            with open(output_path, 'wb') as f:
                f.write(audio_data)

    def get_voice_list(self, setting: Setting = None, force_refresh: bool = False) -> List[Dict[str, Any]]:
        """
        Get the list of available voices from CosyVoice 2
        
        Args:
            setting: Application settings
            force_refresh: If True, force refresh the voice list cache
            
        Returns:
            List of available voices with their properties
        """
        current_time = time.time()
        
        # Return cached voices if available and not expired
        if (not force_refresh and 
            self._voice_list_cache and 
            (current_time - self._voice_list_last_updated) < self._voice_cache_ttl):
            return self._voice_list_cache
            
        try:
            # Make the API request with retry logic
            response = self._make_request(
                method="GET",
                endpoint="/voices",
                params={"details": "true"}  # Request detailed voice information
            )
            
            if response.status_code == 200:
                # Parse the response and normalize the format
                voices_data = response.json()
                
                # Transform the response to a standardized format
                voices = []
                if isinstance(voices_data, dict) and 'voices' in voices_data:
                    # Handle paginated response
                    voices_list = voices_data['voices']
                elif isinstance(voices_data, list):
                    # Handle direct list response
                    voices_list = voices_data
                else:
                    raise ValueError("Unexpected response format from CosyVoice 2 API")
                
                # Process each voice
                for voice in voices_list:
                    try:
                        # Extract language information
                        language_code = voice.get('language_code', '')
                        language_name = voice.get('language_name', '')
                        
                        # Extract styles and emotions if available
                        styles = voice.get('styles', [])
                        emotions = voice.get('emotions', [])
                        
                        # Create a normalized voice entry
                        normalized_voice = {
                            'id': voice.get('id', ''),
                            'name': voice.get('name', 'Unnamed Voice'),
                            'gender': voice.get('gender', 'UNSPECIFIED').title(),
                            'language': language_code,
                            'language_name': language_name,
                            'supported_styles': styles,
                            'supported_emotions': emotions,
                            'provider': 'cosyvoice2',
                            'neural': True,  # CosyVoice 2 is neural by default
                            'sample_rate': voice.get('sample_rate_hertz', self.sample_rate),
                            'description': voice.get('description', '')
                        }
                        
                        voices.append(normalized_voice)
                        
                    except Exception as e:
                        logger.warning(f"Error processing voice {voice.get('id', 'unknown')}: {str(e)}")
                
                # Update cache
                self._voice_list_cache = voices
                self._voice_list_last_updated = current_time
                
                return voices
                
            else:
                logger.warning(f"Failed to get voice list: {response.status_code} - {response.text}")
                return self._get_default_voices()
                
        except Exception as e:
            logger.error(f"Error getting voice list from CosyVoice 2: {str(e)}")
            return self._get_default_voices()
    
    def _get_default_voices(self) -> List[Dict[str, Any]]:
        """Return a default set of voices for fallback"""
        return [
            {
                'id': 'zh-CN-YunxiNeural',
                'name': 'Yunxi (Neural)',
                'gender': 'Male',
                'language': 'zh-CN',
                'language_name': 'Chinese (Mandarin)',
                'supported_styles': ['neutral', 'assistant', 'chat', 'customerservice', 'newscast'],
                'supported_emotions': ['neutral', 'happy', 'sad', 'angry', 'fearful', 'disgust', 'surprised'],
                'provider': 'cosyvoice2',
                'neural': True,
                'sample_rate': 24000,
                'description': 'Chinese (Mandarin) neural voice with multiple styles and emotions'
            },
            {
                'id': 'en-US-JennyNeural',
                'name': 'Jenny (Neural)',
                'gender': 'Female',
                'language': 'en-US',
                'language_name': 'English (United States)',
                'supported_styles': ['neutral', 'assistant', 'chat', 'customerservice', 'newscast'],
                'supported_emotions': ['neutral', 'happy', 'sad', 'angry', 'fearful', 'disgust', 'surprised'],
                'provider': 'cosyvoice2',
                'neural': True,
                'sample_rate': 24000,
                'description': 'English (US) neural voice with multiple styles and emotions'
            },
            {
                'id': 'ja-JP-NanamiNeural',
                'name': 'Nanami (Neural)',
                'gender': 'Female',
                'language': 'ja-JP',
                'language_name': 'Japanese',
                'supported_styles': ['neutral', 'assistant', 'chat', 'customerservice', 'newscast'],
                'supported_emotions': ['neutral', 'happy', 'sad', 'angry', 'fearful', 'disgust', 'surprised'],
                'provider': 'cosyvoice2',
                'neural': True,
                'sample_rate': 24000,
                'description': 'Japanese neural voice with multiple styles and emotions'
            }
        ]

# For testing
if __name__ == "__main__":
    import asyncio
    from dotenv import load_dotenv
    
    async def test():
        load_dotenv()
        service = CosyVoiceTtsService()
        
        # Test TTS
        output_file = "test_cosyvoice.wav"
        await service.tts("这是一个测试，测试CosyVoice语音合成效果。", output_file, None)
        print(f"TTS output saved to {output_file}")
        
        # Test voice list
        voices = service.get_voice_list()
        print("\nAvailable voices:")
        for voice in voices:
            print(f"- {voice['name']} ({voice['id']})")
    
    asyncio.run(test())
